#include "security_util.h"
#include <grpcpp/security/credentials.h>
#include <fstream>
#include <sstream>
#include <iostream>

namespace kwdb {
namespace security {

// Simple Base64 encoding table
static const char base64_chars[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

std::string base64_encode(const std::string& input) {
    std::string ret;
    int i = 0;
    int j = 0;
    unsigned char char_array_3[3];
    unsigned char char_array_4[4];
    size_t in_len = input.length();
    const unsigned char* bytes_to_encode =
        reinterpret_cast<const unsigned char*>(input.c_str());

    while (in_len--) {
        char_array_3[i++] = *(bytes_to_encode++);
        if (i == 3) {
            char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
            char_array_4[1] = ((char_array_3[0] & 0x03) << 4) +
                             ((char_array_3[1] & 0xf0) >> 4);
            char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) +
                             ((char_array_3[2] & 0xc0) >> 6);
            char_array_4[3] = char_array_3[2] & 0x3f;

            for(i = 0; i < 4; i++)
                ret += base64_chars[char_array_4[i]];
            i = 0;
        }
    }

    if (i) {
        for(j = i; j < 3; j++)
            char_array_3[j] = '\0';

        char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
        char_array_4[1] = ((char_array_3[0] & 0x03) << 4) +
                         ((char_array_3[1] & 0xf0) >> 4);
        char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) +
                         ((char_array_3[2] & 0xc0) >> 6);

        for (j = 0; j < i + 1; j++)
            ret += base64_chars[char_array_4[j]];

        while((i++ < 3))
            ret += '=';
    }

    return ret;
}

std::shared_ptr<grpc::ChannelCredentials> create_channel_credentials(
    const SecurityConfig& config) {

    if (!config.tls) {
        // No TLS - use insecure credentials
        return grpc::InsecureChannelCredentials();
    }

    // TLS enabled
    grpc::SslCredentialsOptions ssl_opts;

    if (!config.ca_cert_path.empty()) {
        // Load custom CA certificate
        std::ifstream ca_file(config.ca_cert_path);
        if (!ca_file.is_open()) {
            std::cerr << "[Security] Warning: Failed to open CA cert file: "
                      << config.ca_cert_path << std::endl;
            std::cerr << "[Security] Falling back to system default trust store"
                      << std::endl;
        } else {
            std::stringstream buffer;
            buffer << ca_file.rdbuf();
            ssl_opts.pem_root_certs = buffer.str();
            std::cout << "[Security] Loaded custom CA certificate from: "
                      << config.ca_cert_path << std::endl;
        }
    }

    // If pem_root_certs is empty, gRPC will use system default trust store
    return grpc::SslCredentials(ssl_opts);
}

void add_auth_metadata(grpc::ClientContext& context, const SecurityConfig& config) {
    if (!config.auth_token.empty()) {
        // Bearer token authentication (preferred)
        std::string auth_value = "Bearer " + config.auth_token;
        context.AddMetadata("authorization", auth_value);
        std::cout << "[Security] Using Bearer token authentication" << std::endl;
    }
    else if (!config.username.empty() && !config.password.empty()) {
        // Basic authentication
        std::string credentials = config.username + ":" + config.password;
        std::string encoded = base64_encode(credentials);
        std::string auth_value = "Basic " + encoded;
        context.AddMetadata("authorization", auth_value);
        std::cout << "[Security] Using Basic authentication (user: "
                  << config.username << ")" << std::endl;
    }
}

} // namespace security

// Implementation of SecurityConfig::get_auth_header()
std::string SecurityConfig::get_auth_header() const {
    if (!auth_token.empty()) {
        return "Bearer " + auth_token;
    }
    else if (!username.empty() && !password.empty()) {
        std::string credentials = username + ":" + password;
        std::string encoded = security::base64_encode(credentials);
        return "Basic " + encoded;
    }
    return "";
}

} // namespace kwdb
